from enum import Enum
import numpy as np
import traffic_history_reader
import networkx as nx
import math

class BlockType(Enum):
	SUPERBLOCK = 1
	BORDER_ROUTER = 2

'''
A simple topology class that contains the most basic information
needed for a homogeneous topology
# A network instance must have three aspects, a router 
'''
class AuroraNetwork(object):
	# Initializes the fractional topology instance
	# nblocks - the number of blocks
	# block_names_list - a list of the block names
	# topology engineer - yes or no
	# 					-> if yes then need to define how often and how much for training
	# 
	def __init__(self, dcn_name, block_params, block_names_list, ignore_border_routers=False):
		####################################################
		## A network needs a few key parameters
		## 1) router evaluation
		## 2) path selector
		## 3) traffic engineer
		## 4) topology engineer
		## 5) traffic history
		## 6) other hardware physical parameters, like number of blocks, number of links in every block, etc.
		####################################################
		## Traffic snapshot reader
		## Hardware physical parameters
		self.ignore_border_routers = ignore_border_routers
		self.names_to_id_dict = {}
		self.id_to_names_dict = {}
		self.block_link_capacity = {}
		self.block_num_links = {}
		self.nblocks = 0
		self.dcn_name = dcn_name
		self.flow_id = {}
		self.block_names_list = block_names_list
		
		self._build_network_physical_properties(block_names_list, block_params, ignore_border_routers)
		self.total_links = 0
		for block in range(self.nblocks):
			self.total_links += self.block_num_links[block]
		index = 0
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				if i != j:
					self.flow_id[(i,j,)] = index
					index += 1
		print("Physical Network Initialized")
		return	

	def get_dcn_name(self):
		return self.dcn_name

	def get_total_links(self):
		return self.total_links
	'''
	Building functions, including the 1) physical network parameters 
									  2) topology engineer 
									  3) traffic engineer
	'''
	# internal method used for building the entire network
	def _build_network_physical_properties(self, block_names_list, block_params, ignore_border_routers):
		print("Building Aurora Topology's Physical Hardware")
		superblock_id_offset = 0
		borderrouter_id_offset = 0
		for block_name in block_names_list:
			if "ju" in block_name:
				self.names_to_id_dict[block_name] = superblock_id_offset
				superblock_id_offset += 1
			elif "ExternalBlock" in block_name and not self.ignore_border_routers:
				self.names_to_id_dict[block_name] = borderrouter_id_offset
				borderrouter_id_offset += 1
			else:
				raise Exception("Unrecognizeable Block Name: {}".format(block_name))
		for block_name in block_names_list:
			if "ju" in block_name:
				superblock_id = self.names_to_id_dict[block_name]
				self.id_to_names_dict[superblock_id] = block_name
				self.block_link_capacity[superblock_id] = block_params[BlockType.SUPERBLOCK]["link capacity"]
				self.block_num_links[superblock_id] = block_params[BlockType.SUPERBLOCK]["num links"]
			elif "ExternalBlock" in block_name and not self.ignore_border_routers:
				border_id = superblock_id_offset + self.names_to_id_dict[block_name]
				self.id_to_names_dict[border_id] = block_name
				self.names_to_id_dict[block_name] = border_id
				self.block_link_capacity[border_id] = block_params[BlockType.BORDER_ROUTER]["link capacity"]
				self.block_num_links[border_id] = block_params[BlockType.BORDER_ROUTER]["num links"]
			else:
				raise Exception("Unrecognizeable Block Name: {}".format(block_name))
		if self.ignore_border_routers:
			self.nblocks = superblock_id_offset
		else:
			self.nblocks = superblock_id_offset + borderrouter_id_offset
		print("Aurora Topology's Physical Hardware Building COMPLETED")
		return 

	def get_link_capacity(self, block_id):
		return self.block_link_capacity[block_id]

	def get_num_links(self, block_id):
		return self.block_num_links[block_id]
	
	def get_num_blocks(self):
		return self.nblocks

	def get_traffic_vector_from_traffic_snapshot(self, traffic_snapshot, normalize_to=-1.):
		tm_snapshot_point = np.zeros(self.nblocks * (self.nblocks - 1)) 
		traffic_sum = 0.
		for record in traffic_snapshot.records:
			src_string = record.source
			dst_string = record.destination
			if src_string == dst_string:
				continue
			numbytes = float(record.max_load_gbps)
			src_id = None
			dst_id = None
			if src_string not in self.names_to_id_dict or dst_string not in self.names_to_id_dict:
				continue
			else:
				src_id = self.names_to_id_dict[src_string]
				dst_id = self.names_to_id_dict[dst_string]
			traffic_sum += numbytes
			index = self.flow_id[(src_id, dst_id,)]
			tm_snapshot_point[index] += numbytes
		if normalize_to > 0:
			tm_snapshot_point = [(x / traffic_sum) * normalize_to for x in tm_snapshot_point]
		return tm_snapshot_point

	def get_traffic_matrix_from_traffic_snapshot(self, traffic_snapshot, normalize_to=-1.):
		tm = np.zeros((self.nblocks, self.nblocks))
		traffic_sum = 0.
		for record in traffic_snapshot.records:
			src_string = record.source
			dst_string = record.destination
			if src_string == dst_string:
				continue
			numbytes = float(record.max_load_gbps)
			src_id = None
			dst_id = None
			if src_string not in self.names_to_id_dict or dst_string not in self.names_to_id_dict:
				continue
			else:
				src_id = self.names_to_id_dict[src_string]
				dst_id = self.names_to_id_dict[dst_string]
			traffic_sum += numbytes
			tm[src_id][dst_id] += numbytes
		if normalize_to > 0:
			for i in range(self.nblocks):
				for j in range(self.nblocks):
					if i != j:
						tm[src_id][dst_id] = tm[src_id][dst_id] * normalize_to / traffic_sum
		return tm

	def transform_traffic_matrix_to_traffic_point(self, traffic_matrix):
		traffic_point = np.zeros((self.nblocks * (self.nblocks - 1), ))
		index = 0
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				if i != j:
					traffic_point[index] = traffic_matrix[i][j]
					index += 1
		return traffic_point

	def transform_traffic_point_to_traffic_matrix(self, traffic_point):
		dim = self.nblocks * (self.nblocks - 1)
		assert(len(traffic_point) == dim)
		tm = np.zeros((self.nblocks, self.nblocks,))
		index = 0
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				if i != j:
					tm[i][j] = traffic_point[index]
					index += 1
		return tm

	def normalize_matrix(self, matrix, norm=1.):
		matrix_sum = 0.
		nentry = len(matrix)
		for i in range(nentry):
			for j in range(nentry):
				matrix_sum += float(matrix[i][j])
		normalized_matrix = np.zeros((nentry, nentry,))
		for i in range(nentry):
			for j in range(nentry):
				normalized_matrix[i][j] = (matrix[i][j] / matrix_sum) * norm
		return normalized_matrix

	def normalize_point(self, point, norm=1.):
		point_sum = sum(point)
		nentry = len(point)
		normalized_point = np.zeros((nentry,))
		for i in range(nentry):
			normalized_point[i] = (point[i] / matrix_sum) * norm
		return normalized_point		

	def print_stats(self, adj_matrix):
		nblocks = self.nblocks
		print("")
		for pod in range(self.nblocks):
			row_sum = 0.
			col_sum = 0.
			for target_pod in range(self.nblocks):
				if pod != target_pod:
					row_sum += adj_matrix[pod][target_pod]
					col_sum += adj_matrix[target_pod][pod]
			print("pod {}, rowsum: {}, colsum: {}".format(pod, row_sum, col_sum))
		print("")
		return

	def get_traffic_history(self):
		return self.traffic_history

	## assumes that the entire topology is connected via a giant switch
	def round_fractional_topology_giant_switch(self, fractional_adj_matrix, pairs_to_ignore, rounding_simple=False):
		integer_adj_matrix = np.zeros((self.nblocks, self.nblocks))
		nnodes = 2 + (2 * self.nblocks)
		G = nx.DiGraph()
		edges = []
		# add edges from src to first layer nodes and between second layer nodes to sink
		for i in range(self.nblocks):
			egress_sum = 0.
			ingress_sum = 0.
			for j in range(self.nblocks):
				if i != j:
					egress_sum += math.floor(fractional_adj_matrix[i][j])
					ingress_sum += math.floor(fractional_adj_matrix[j][i])
			edges.append((0, i + 1, {'capacity' : self.block_num_links[i] - egress_sum, 'weight': 0}))
			edges.append((self.nblocks + i + 1, nnodes - 1, {'capacity' : self.block_num_links[i] - ingress_sum, 'weight' : 0}))
		# next, add edges between first layer nodes and second layer nodes
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				#if i != j and (i,j) not in pairs_to_ignore:
				if i != j:
					score = int((math.ceil(fractional_adj_matrix[i][j]) - fractional_adj_matrix[i][j]) * 1E6)
					edges.append((i + 1, self.nblocks + j + 1, {'capacity' : 1, 'weight': score}))

		# next, add the edges set into the graph
		G.add_edges_from(edges)
		mincostFlow = nx.max_flow_min_cost(G, 0, nnodes - 1)
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				if i != j and (i,j) not in pairs_to_ignore:
					#assert(mincostFlow[i + 1][len(ocs_striping) + j + 1] >= 0)
					assert(mincostFlow[i + 1][self.nblocks + j + 1] <= 1)
					integer_adj_matrix[i][j] = math.floor(fractional_adj_matrix[i][j]) + mincostFlow[i + 1][self.nblocks + j + 1]
		## Finally, check to see if some links are 0
		for i in range(self.nblocks):
			for j in range(self.nblocks):
				if i != j and (i,j) not in pairs_to_ignore and integer_adj_matrix[i][j] < 1:
					if rounding_simple:
						integer_adj_matrix[i][j] = 1.
					else:
						i_prime = -1
						max_iprime = -1
						for ii in range(self.nblocks):
							if ii != j and integer_adj_matrix[ii][j] > max_iprime and integer_adj_matrix[ii][j] > 1:
								max_iprime = integer_adj_matrix[ii][j]
								i_prime = ii
						j_prime = -1
						max_jprime = -1
						for jj in range(self.nblocks):
							if jj != i and integer_adj_matrix[i][jj] > max_jprime and integer_adj_matrix[i][jj] > 1:
								max_jprime = integer_adj_matrix[i][jj]
								j_prime = jj
						if i_prime < 0 or j_prime < 0:
							continue
						#assert(i_prime >= 0 and j_prime >= 0)
						integer_adj_matrix[i][j_prime] -= 1
						integer_adj_matrix[i_prime][j] -= 1
						integer_adj_matrix[i][j] += 1
						integer_adj_matrix[i_prime][j_prime] += 1
		return integer_adj_matrix
